# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# Modified by Sen Yang (yangsenius@seu.edu.cn)
# ------------------------------------------------------------------------------

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import os
import pprint
import shutil
import sys

import numpy as np
import random

import torch
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
from tensorboardX import SummaryWriter

# 这个模块在sys.path中添加了路径G:\TransPose-main\tools\..\lib 和 G:\TransPose-main\tools\..\lib/poseeval/py-motmetrics
import _init_paths
# 由于添加了路径，所以下面导入路径的方法是正确的
from config import cfg
from config import update_config
from core.loss import JointsMSELoss
from core.function import train
from core.function import validate
from utils.utils import get_optimizer
from utils.utils import save_checkpoint
from utils.utils import create_logger
from utils.utils import get_model_summary

# from lib.config import cfg
# from lib.config import update_config
# from lib.core.loss import JointsMSELoss
# from lib.core.function import train
# from lib.core.function import validate
# from lib.utils.utils import get_optimizer
# from lib.utils.utils import save_checkpoint
# from lib.utils.utils import create_logger
# from lib.utils.utils import get_model_summary
# TODO 暂时注释掉
# model = eval('models.' + cfg.MODEL.NAME + '.get_pose_net')(
#     cfg, is_train=True
# )
import dataset
import models


def parse_args():
    parser = argparse.ArgumentParser(description='Train keypoints network')  # 创建解析器
    # general
    # 添加参数
    '''
    ArgumentParser.add_argument(name or flags...[, action][, nargs][, const][, default][, type][, choices][, required][, help][, metavar][, dest])

    name or flags - 一个命名或者一个选项字符串的列表，例如 foo 或 -f, --foo。
    action - 当参数在命令行中出现时使用的动作基本类型。
    nargs - 命令行参数应当消耗的数目。
    const - 被一些 action 和 nargs 选择所需求的常数。
    default - 当参数未在命令行中出现时使用的值。
    type - 命令行参数应当被转换成的类型。
    choices - 可用的参数的容器。
    required - 此命令行选项是否可省略 （仅选项可用）。
    help - 一个此选项作用的简单描述。
    metavar - 在使用方法消息中使用的参数值示例。
    dest - 被添加到 parse_args() 所返回对象上的属性名。
    '''
    # 在命令行使用相应的命令 后加该命令的值，就是在赋值了
    parser.add_argument('--cfg',
                        help='experiment configure file name',
                        required=True,
                        type=str)

    parser.add_argument('opts',
                        help="Modify config options using the command-line",
                        default=None,
                        nargs=argparse.REMAINDER)

    # philly
    parser.add_argument('--modelDir',
                        help='model directory',
                        type=str,
                        default='')
    parser.add_argument('--logDir',
                        help='log directory',
                        type=str,
                        default='')
    parser.add_argument('--dataDir',
                        help='data directory',
                        type=str,
                        default='')
    parser.add_argument('--prevModelDir',
                        help='prev Model directory',
                        type=str,
                        default='')

    args = parser.parse_args()

    return args


def main():
    args = parse_args()
    # 这个就是来更新参数，把args里面的参数替换cfg里面设置的参数值
    update_config(cfg, args)  # from lib.config import cfg ==》_C

    # 创建训练的日志
    logger, final_output_dir, tb_log_dir = create_logger(
        cfg, args.cfg, 'train')

    logger.info(pprint.pformat(args))  # 打印信息
    logger.info(cfg)

    '''
    使用非确定性算法，可以增加运行效率
    如果设置为torch.backends.cudnn.enabled =True，说明设置为使用使用非确定性算法
    然后再设置：
    torch.backends.cudnn.benchmark = true
    那么cuDNN使用的非确定性算法就会自动寻找最适合当前配置的高效算法，来达到优化运行效率的问题
    torch.backends.cudnn.deterministic
    如果deterministic为true的话，每次返回的卷积的算法将是确定的（固定的）。
    如果为False的话，每次返回的算法是不固定的
    '''
    # cudnn related setting
    cudnn.benchmark = cfg.CUDNN.BENCHMARK
    torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
    torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED

    '''
    seed
    保证每次实验生成的随机数的值的顺序是一样
    import torch
    seed = 22
    torch.manual_seed(seed)
    Out[4]: <torch._C.Generator at 0x1d408b6dd98>
    torch.rand(1)
    Out[5]: tensor([0.3659])
    torch.rand(1)
    Out[6]: tensor([0.7025])
    退出程序，下次重新执行这个文件的顺序也是一样的值
    '''
    seed = 22
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

    # cfg.MODEL.NAME='pose_hrnet'
    # eval用来执行一个字符串表达式，并返回表达式的值
    # TODO 这个路径可能对不上cfg.MODEL.NAME  《==》  models.transpose_h.py
    model = eval('models.' + cfg.MODEL.NAME + '.get_pose_net')(
        cfg, is_train=True
    )

    # copy model file
    '''
    __file__是内置的变量，保存文件的绝对路径，F:/研究生阶段学习/骆老师研究生团队/论文复现/TransPose-main/tools/train.py
    os.path.dirname(__file__)获取除了文件名以外的绝对路径 F: / 研究生阶段学习 / 骆老师研究生团队 / 论文复现 / TransPose - main / tools
    
    os.path.join(this_dir, '../lib/models', cfg.MODEL.NAME + '.py'):
    F:/研究生阶段学习/骆老师研究生团队/论文复现/TransPose-main/tools==>
    F:/研究生阶段学习/骆老师研究生团队/论文复现/TransPose-main/==>
    F:/研究生阶段学习/骆老师研究生团队/论文复现/TransPose-main/lib/models==>
    F:/研究生阶段学习/骆老师研究生团队/论文复现/TransPose-main/lib/models/transpose.py
    
    shutil 是对一系列文件(或者文件集合)进行操作，不是单个文件
    shutil.copy2(src,dest)  将src文件拷贝到dest目录下面
    '''
    this_dir = os.path.dirname(__file__)
    shutil.copy2(
        os.path.join(this_dir, '../lib/models', cfg.MODEL.NAME + '.py'),
        final_output_dir)
    # logger.info(pprint.pformat(model))

    '''
    tensorboardX可视化工具
    首先是：定义一个实例SummerWriter(),参数为：
    def __init__(self, log_dir=None, comment='', **kwargs):
     其中log_dir为生成的文件所放的目录，comment为文件名称。默认目录为生成runs文件夹目录。
    '''
    writer_dict = {
        'writer': SummaryWriter(log_dir=tb_log_dir),
        'train_global_steps': 0,
        'valid_global_steps': 0,
    }

    # TODO 不理解这个是干什么的
    dump_input = torch.rand(
        (1, 3, cfg.MODEL.IMAGE_SIZE[1], cfg.MODEL.IMAGE_SIZE[0])
    )

    model = torch.nn.DataParallel(model, device_ids=[0,]).cuda()  # _C.GPUS = (0,)是一个元组类型的数据

    # define loss function (criterion) and optimizer
    # 自定义loss
    criterion = JointsMSELoss(
        use_target_weight=cfg.LOSS.USE_TARGET_WEIGHT
    ).cuda()

    # Data loading code
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
    )
    train_dataset = eval('dataset.' + cfg.DATASET.DATASET)(
        cfg, cfg.DATASET.ROOT, cfg.DATASET.TRAIN_SET, True,
        transforms.Compose([
            transforms.ToTensor(),
            normalize,
        ])
    )
    valid_dataset = eval('dataset.' + cfg.DATASET.DATASET)(
        cfg, cfg.DATASET.ROOT, cfg.DATASET.TEST_SET, False,
        transforms.Compose([
            transforms.ToTensor(),
            normalize,
        ])
    )

    # batch_size=cfg.TRAIN.BATCH_SIZE_PER_GPU * len(cfg.GPUS),
    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=96,
        shuffle=cfg.TRAIN.SHUFFLE,
        num_workers=0,
        pin_memory=cfg.PIN_MEMORY
    )
    # batch_size=cfg.TEST.BATCH_SIZE_PER_GPU * len(cfg.GPUS)
    valid_loader = torch.utils.data.DataLoader(
        valid_dataset,
        batch_size=96,
        shuffle=False,
        num_workers=0,
        pin_memory=cfg.PIN_MEMORY
    )

    best_perf = 0.0
    best_model = False
    last_epoch = -1

    optimizer = get_optimizer(cfg, model)

    begin_epoch = cfg.TRAIN.BEGIN_EPOCH
    checkpoint_file = os.path.join(
        final_output_dir, 'checkpoint.pth'
    )

    if cfg.AUTO_RESUME and os.path.exists(checkpoint_file):
        logger.info("=> loading checkpoint '{}'".format(checkpoint_file))
        checkpoint = torch.load(checkpoint_file)
        begin_epoch = checkpoint['epoch']
        best_perf = checkpoint['perf']
        last_epoch = checkpoint['epoch']

        writer_dict['train_global_steps'] = checkpoint['train_global_steps']
        writer_dict['valid_global_steps'] = checkpoint['valid_global_steps']

        model.load_state_dict(checkpoint['state_dict'])

        optimizer.load_state_dict(checkpoint['optimizer'])
        logger.info("=> loaded checkpoint '{}' (epoch {})".format(
            checkpoint_file, checkpoint['epoch']))

    # lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
    #     optimizer, cfg.TRAIN.LR_STEP, cfg.TRAIN.LR_FACTOR,
    #     last_epoch=last_epoch
    # )

    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, cfg.TRAIN.END_EPOCH, eta_min=cfg.TRAIN.LR_END, last_epoch=last_epoch)

    model.cuda()

    for epoch in range(begin_epoch, cfg.TRAIN.END_EPOCH):

        logger.info("=> current learning rate is {:.6f}".format(lr_scheduler.get_last_lr()[0]))
        # train for one epoch
        train(cfg, train_loader, model, criterion, optimizer, epoch,
              final_output_dir, tb_log_dir, writer_dict)

        # evaluate on validation set
        perf_indicator = validate(
            cfg, valid_loader, valid_dataset, model, criterion,
            final_output_dir, tb_log_dir, writer_dict
        )

        lr_scheduler.step()

        if perf_indicator >= best_perf:
            best_perf = perf_indicator
            best_model = True
        else:
            best_model = False

        logger.info('=> saving checkpoint to {}'.format(final_output_dir))
        save_checkpoint({
            'epoch': epoch + 1,
            'model': cfg.MODEL.NAME,
            'state_dict': model.state_dict(),
            'best_state_dict': model.module.state_dict(),
            'perf': perf_indicator,
            'optimizer': optimizer.state_dict(),
            'train_global_steps': writer_dict['train_global_steps'],
            'valid_global_steps': writer_dict['valid_global_steps'],
        }, best_model, final_output_dir)

    final_model_state_file = os.path.join(
        final_output_dir, 'final_state.pth'
    )
    cache_model_state_file = os.path.join(
        '/cache', 'final_state.pth'
    )
    logger.info('=> saving final model state to {}'.format(
        final_model_state_file)
    )
    torch.save(model.module.state_dict(), final_model_state_file)
    torch.save(model.module.state_dict(), cache_model_state_file)
    writer_dict['writer'].close()


if __name__ == '__main__':
    main()
    # print(type(cfg))  # cfg参数的值再lib.default这个文件里面找，还有就是再配置文件里面找
    # print(cfg['MODEL']['INIT_WEIGHTS'])  # 可以使用字典的形式提取值
    # print(cfg['MODEL']['PRETRAINED'])  # 这个值一开始是空的
    # print(__file__)  # __file__是内置的变量，保存文件的绝对路径
    # print(os.path.dirname(__file__))  # 除了文件名以外的绝对路径
    # print(sys.modules)  # 模块的搜索路径
